package com.rmd.ai.grpc.handler;

import com.google.protobuf.ByteString;
import com.rmd.ai.grpc.service.Ability;
import com.rmd.ai.grpc.service.AbilityGrpcChannelStatePool;
import com.rmd.ai.grpc.service.RmdAiServiceGrpc;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.stub.StreamObserver;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.util.List;
import java.util.concurrent.TimeUnit;

@Slf4j
@Component
public class AbilityGrpcHandler {
    private ManagedChannel channel;
    private RmdAiServiceGrpc.RmdAiServiceStub stub;
    @Value("${grpc.client.ability.server}")
    private String aiServer;


    @Autowired
    private AbilityGrpcChannelStatePool channelStatePool;

    @PostConstruct
    public void init () {
        String[] ips = aiServer.split(":");
        channel = ManagedChannelBuilder.forAddress(ips[0], Integer.parseInt(ips[1]))
                .usePlaintext()
                .build();//池化处理 成本高
        stub = RmdAiServiceGrpc.newStub(channel);

    }

    public void shutdown() {
        try {
            channel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            log.error("",e);
        }
    }


    public void asr(Long appId,String appSecret, Integer fs, ByteString audioData, final AbilityGrpcCallback callBack) {
        StreamObserver<Ability.RmdAsrResponse> responseObserver = new StreamObserver<Ability.RmdAsrResponse>() {
            public void onNext(Ability.RmdAsrResponse result) {
                //log.info("grpc asr onNext:{}",result.getResult());
                //实例方法回调
                callBack.callback(result.getResult());
            }

            public void onError(Throwable throwable) {
                log.error("grpc asr onError",throwable);
            }

            public void onCompleted() {
                log.debug("grpc asr onCompleted");
                callBack.completed();
            }

        };
        Ability.RmdAsrRequest request = Ability.RmdAsrRequest.newBuilder().setAppId(appId).setAppSecret(appSecret).setFs(fs).setAudioData(audioData).build();
        try {
            stub.asr(request,  responseObserver);
        }catch (Exception e){
            log.error("asr ",e);
        }
    }
    public void asrLarge(Long appId,String appSecret, Integer fs, ByteString audioData, final AbilityGrpcCallback callBack) {
        StreamObserver<Ability.RmdAsrResponse> responseObserver = new StreamObserver<Ability.RmdAsrResponse>() {
            public void onNext(Ability.RmdAsrResponse result) {
                log.info("grpc asr onNext:{}",result.getResult());
                //实例方法回调
                callBack.callback(result.getResult());
            }

            public void onError(Throwable throwable) {
                log.error("grpc asr onError",throwable);
            }

            public void onCompleted() {
                log.debug("grpc asr onCompleted");
            }

        };
        Ability.RmdAsrRequest request = Ability.RmdAsrRequest.newBuilder().setAppId(appId).setAppSecret(appSecret).setFs(fs).setAudioData(audioData).build();
        try {
            stub.asrLarge(request,  responseObserver);
        }catch (Exception e){
            log.error("asr ",e);
        }
    }
    public void embedding(Long appId,String appSecret, String text, final AbilityGrpcCallback callBack) {
        StreamObserver<Ability.RmdNlpResponse> responseObserver = new StreamObserver<Ability.RmdNlpResponse>() {
            public void onNext(Ability.RmdNlpResponse result) {
                log.debug("grpc embedding onNext:{}",result.getResult());
                //实例方法回调
                callBack.callback(result.getResult());
            }

            public void onError(Throwable throwable) {
                log.error("grpc asr onError",throwable);
            }

            public void onCompleted() {
                log.debug("grpc asr onCompleted");
            }

        };
        Ability.RmdEmbeddingRequest request = Ability.RmdEmbeddingRequest.newBuilder().setAppId(appId).setAppSecret(appSecret).setText(text).build();
        try {
            stub.embedding(request,  responseObserver);
        }catch (Exception e){
            log.error("embedding: ",e);
        }
    }
    public void ranking(Long appId,String appSecret, String text, List<Ability.RmdFaqSupportSet> supportSet, Double minScore, final AbilityGrpcCallback callBack) {
        StreamObserver<Ability.RmdNlpResponse> responseObserver = new StreamObserver<Ability.RmdNlpResponse>() {
            public void onNext(Ability.RmdNlpResponse result) {
                log.debug("grpc ranking onNext:{}",result.getResult());
                //实例方法回调
                callBack.callback(result.getResult());
            }

            public void onError(Throwable throwable) {
                log.error("grpc ranking onError",throwable);
            }

            public void onCompleted() {
                log.debug("grpc ranking onCompleted");
            }

        };

        Ability.RmdFaqRequest request = Ability.RmdFaqRequest.newBuilder().setAppId(appId).setAppSecret(appSecret).setQuery(text).addAllSupportSet(supportSet).setMinScore(minScore == null ? 0 : minScore).build();
        try {
            stub.faq(request,  responseObserver);
        }catch (Exception e){
            log.error("ranking",e);
        }
    }

    public void rankingBatch(Long appId,String appSecret, List<Ability.RmdFaqBatchItem> faqItemList, Double minScore, final AbilityGrpcCallback callBack) {
        StreamObserver<Ability.RmdNlpResponse> responseObserver = new StreamObserver<Ability.RmdNlpResponse>() {
            public void onNext(Ability.RmdNlpResponse result) {
                log.debug("grpc ranking batch onNext:{}",result.getResult());
                //实例方法回调
                callBack.callback(result.getResult());
            }

            public void onError(Throwable throwable) {
                log.error("grpc ranking batch onError",throwable);
            }

            public void onCompleted() {
                log.debug("grpc ranking batch onCompleted");
            }

        };

        Ability.RmdFaqBatchRequest request = Ability.RmdFaqBatchRequest.newBuilder().setAppId(appId).setAppSecret(appSecret).addAllItems(faqItemList).setMinScore(minScore == null ? 0 : minScore).build();
        try {
            stub.faqBatch(request,  responseObserver);
        }catch (Exception e){
            log.error("ranking batch",e);
        }
    }


    public void tts(Long appId,String appSecret,Float scale, Integer fs, String text,String voice,final AbilityGrpcCallback callBack) {
        StreamObserver<Ability.RmdTtsResponse> responseObserver = new StreamObserver<Ability.RmdTtsResponse>() {
            public void onNext(Ability.RmdTtsResponse result) {
                log.debug("grpc tts onNext ");
                //实例方法回调
                callBack.callback(result.getAudioData().toByteArray());
            }

            public void onError(Throwable throwable) {
                log.error("grpc tts onError",throwable);
                callBack.completed();
            }

            public void onCompleted() {
                log.debug("grpc tts onCompleted");
                callBack.completed();
            }

        };
        Ability.RmdTtsRequest request = Ability.RmdTtsRequest.newBuilder().setAppId(appId).setAppSecret(appSecret).setScale(scale).setFs(fs).setText(text).setVoice(voice).build();
        try {
            stub.tts(request,  responseObserver);
        }catch (Exception e){
            log.error("tts" ,e);
        }
    }

    public void nlu(Long appId,String appSecret, String text,String schema, final AbilityGrpcCallback callBack) {
        StreamObserver<Ability.RmdNlpResponse> responseObserver = new StreamObserver<Ability.RmdNlpResponse>() {
            public void onNext(Ability.RmdNlpResponse result) {
                log.debug("grpc nlu onNext:{}",result.getResult());
                //实例方法回调
                callBack.callback(result.getResult());
            }

            public void onError(Throwable throwable) {
                log.error("grpc nlu onError",throwable);
            }

            public void onCompleted() {
                log.debug("grpc nlu onCompleted");
            }

        };
        Ability.RmdNluRequest request = Ability.RmdNluRequest.newBuilder().setAppId(appId).setAppSecret(appSecret).setText(text).setNluSchema(schema).build();
        try {
            stub.nlu(request,  responseObserver);
        }catch (Exception e){
            log.error("nlu" ,e);
        }
    }
    public void asrPipelineForNav(Long appId,String appSecret, Integer fs, ByteString audioData,String schema, final AbilityGrpcCallback callBack) {
        StreamObserver<Ability.RmdPipelineResponse> responseObserver = new StreamObserver<Ability.RmdPipelineResponse>() {
            public void onNext(Ability.RmdPipelineResponse result) {
                log.info("grpc pipeline onNext:{}",result.getResult());
                //实例方法回调
                callBack.callback(result.getResult());
            }

            public void onError(Throwable throwable) {

                log.error("grpc pipeline onError",throwable);
            }

            public void onCompleted() {
                log.debug("grpc pipeline onCompleted");
            }

        };
        Ability.RmdPipelineRequest request = Ability.RmdPipelineRequest.newBuilder()
                .setAppId(appId).setAppSecret(appSecret)
                .setAsr(true).setFs(fs).setAudioData(audioData)
                .setStopword(true).setTextCorrect(true).setItn(true)
                .setEmbedding(true)
                .setNlu(true).setNluSchema(schema)
                .setRewrite(true).setRewriteEmbedding(true).build();
        try {
            stub.pipeline(request,  responseObserver);
        }catch (Exception e){
            log.error("pipeline" ,e);
        }
    }
    public void asrForPool(Long appId,String appSecret, Integer fs, ByteString audioData, final AbilityGrpcCallback callBack) {
        ManagedChannel channel = null;
        try {
            channel = channelStatePool.borrowConnect();
            RmdAiServiceGrpc.RmdAiServiceStub stub = RmdAiServiceGrpc.newStub(channel);
            ManagedChannel finalChannel = channel;
            StreamObserver<Ability.RmdAsrResponse> responseObserver = new StreamObserver<Ability.RmdAsrResponse>() {
                public void onNext(Ability.RmdAsrResponse result) {
                    //log.info("grpc asr onNext:{}",result.getResult());
                    //实例方法回调
                    callBack.callback(result.getResult());
                }

                public void onError(Throwable throwable) {
                    log.error("grpc asr onError",throwable);
                    // 归还连接对象
                    if (finalChannel != null) {
                        channelStatePool.returnConnect(finalChannel);
                    }
                }

                public void onCompleted() {
                    log.debug("grpc asr onCompleted");
                    callBack.completed();
                    // 归还连接对象
                    if (finalChannel != null) {
                        channelStatePool.returnConnect(finalChannel);
                    }
                }

            };
            Ability.RmdAsrRequest request = Ability.RmdAsrRequest.newBuilder().setAppId(appId).setAppSecret(appSecret).setFs(fs).setAudioData(audioData).build();

            stub.asr(request,  responseObserver);
        }catch (Exception e){
            log.error("asr for pool",e);
            // 归还连接对象
            if (channel != null) {
                channelStatePool.returnConnect(channel);
            }
        }
    }
}
